class Solution
{
public:
    int largestValsFromLabels(vector<int> &values, vector<int> &labels, int numWanted, int useLimit)
    {
        int n = values.size();
        vector<int> id(n);
        iota(id.begin(), id.end(), 0);
        sort(id.begin(), id.end(), [&](int i, int j)
             { return values[i] > values[j]; });

        int result = 0;
        int selectCount = 0;
        unordered_map<int, int> labelCount;
        for (int i : id)
        {
            if (labelCount[labels[i]] == useLimit)
            {
                continue;
            }
            result += values[i];
            ++labelCount[labels[i]];
            ++selectCount;
            if (selectCount == numWanted)
            {
                break;
            }
        }
        return result;
    }
};